#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Data models for the visualization engine

This module defines the core data structures used throughout the engine.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Any, Optional, Union
import pandas as pd
from datetime import datetime


@dataclass
class Dataset:
    """数据集模型"""
    df: pd.DataFrame
    dataset_description: str
    table_description: str
    field_descriptions: Dict[str, str]
    dataset_id: str = ""
    sample_data: Optional[pd.DataFrame] = None
    created_at: datetime = field(default_factory=datetime.now)
    
    def __post_init__(self):
        """初始化后处理"""
        # 清理列名，移除空格和BOM字符
        self.df.columns = self.df.columns.str.strip()
        self.df.columns = self.df.columns.str.replace('\ufeff', '')
        
        if self.sample_data is None:
            self.sample_data = self.df.head(3)
        
        if not self.dataset_id:
            self.dataset_id = f"dataset_{int(self.created_at.timestamp())}"
    
    def get_field_info(self, field_name: str) -> Dict[str, Any]:
        """获取字段信息"""
        if field_name not in self.df.columns:
            raise ValueError(f"字段 '{field_name}' 不存在于数据集中")
        
        series = self.df[field_name]
        return {
            "name": field_name,
            "description": self.field_descriptions.get(field_name, ""),
            "dtype": str(series.dtype),
            "null_count": series.isnull().sum(),
            "unique_count": series.nunique(),
            "sample_values": series.dropna().head(5).tolist()
        }
    
    def get_numeric_fields(self) -> List[str]:
        """获取数值型字段"""
        return self.df.select_dtypes(include=['number']).columns.tolist()
    
    def get_categorical_fields(self) -> List[str]:
        """获取分类型字段"""
        return self.df.select_dtypes(include=['object', 'category']).columns.tolist()
    
    def get_datetime_fields(self) -> List[str]:
        """获取日期时间字段"""
        return self.df.select_dtypes(include=['datetime']).columns.tolist()
    
    def validate(self) -> Dict[str, Any]:
        """验证数据集完整性"""
        issues = []
        
        # 检查DataFrame是否为空
        if self.df.empty:
            issues.append("数据集为空")
        
        # 检查字段描述完整性
        missing_descriptions = []
        for col in self.df.columns:
            if col not in self.field_descriptions:
                missing_descriptions.append(col)
        
        if missing_descriptions:
            issues.append(f"缺少字段描述: {missing_descriptions}")
        
        # 检查数据质量
        null_counts = self.df.isnull().sum()
        high_null_fields = null_counts[null_counts > len(self.df) * 0.5].index.tolist()
        
        if high_null_fields:
            issues.append(f"字段空值过多(>50%): {high_null_fields}")
        
        return {
            "is_valid": len(issues) == 0,
            "issues": issues,
            "row_count": len(self.df),
            "column_count": len(self.df.columns),
            "memory_usage": self.df.memory_usage(deep=True).sum()
        }


@dataclass
class ChartAnalysis:
    """图表分析结果模型"""
    chart_type: str
    need_chart: bool
    x_field: Optional[str] = None
    y_field: Optional[str] = None
    y_fields: List[str] = field(default_factory=list)  # 支持多个y字段
    name_field: Optional[str] = None
    value_field: Optional[str] = None
    group_by: Optional[str] = None
    filter_condition: Optional[str] = None
    aggregation: Optional[str] = None
    sort_by: Optional[str] = None
    sort_order: str = "asc"
    limit: Optional[int] = None
    confidence: float = 0.0
    reasoning: str = ""
    series_config: Dict[str, Any] = field(default_factory=dict)  # 序列配置
    # 执行顺序相关字段
    execution_order: List[str] = field(default_factory=list)  # 执行步骤顺序
    
    def __post_init__(self):
        """验证分析结果"""
        if not self.chart_type:
            raise ValueError("图表类型不能为空")
        
        if self.confidence < 0 or self.confidence > 1:
            raise ValueError("置信度必须在0-1之间")
    
    def to_dict(self, include_debug_info: bool = True) -> Dict[str, Any]:
        """转换为字典格式"""
        result = {
            "chart_type": self.chart_type,
            "need_chart": self.need_chart,
            "x_field": self.x_field,
            "y_field": self.y_field,
            "y_fields": self.y_fields,
            "name_field": self.name_field,
            "value_field": self.value_field,
            "group_by": self.group_by,
            "filter_condition": self.filter_condition,
            "aggregation": self.aggregation,
            "sort_by": self.sort_by,
            "sort_order": self.sort_order,
            "limit": self.limit,
            "confidence": self.confidence,
            "series_config": self.series_config,
            "execution_order": self.execution_order
        }
        
        # 只在debug模式下包含reasoning
        if include_debug_info:
            result["reasoning"] = self.reasoning
            
        return result


@dataclass
class ProcessingResult:
    """数据处理结果模型"""
    processed_df: pd.DataFrame
    processing_code: str
    execution_time: float
    success: bool
    error_message: str = ""
    warnings: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def __post_init__(self):
        """验证处理结果"""
        if not self.success and not self.error_message:
            raise ValueError("处理失败时必须提供错误信息")
    
    def get_summary(self) -> Dict[str, Any]:
        """获取处理结果摘要"""
        return {
            "success": self.success,
            "execution_time": self.execution_time,
            "row_count": len(self.processed_df) if self.success else 0,
            "column_count": len(self.processed_df.columns) if self.success else 0,
            "has_warnings": len(self.warnings) > 0,
            "warning_count": len(self.warnings),
            "error_message": self.error_message
        }


@dataclass
class VisualizationResult:
    """可视化结果模型"""
    result_type: str  # "echarts" or "table"
    config: Dict[str, Any]
    chart_analysis: ChartAnalysis
    processing_result: ProcessingResult
    generation_time: float
    success: bool
    error_message: str = ""
    
    def __post_init__(self):
        """验证可视化结果"""
        if self.result_type not in ["echarts", "table"]:
            raise ValueError("结果类型必须是 'echarts' 或 'table'")
        
        if not self.success and not self.error_message:
            raise ValueError("生成失败时必须提供错误信息")
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        result = {
            "type": self.result_type,
            "success": self.success,
            "generation_time": self.generation_time
        }
        
        if self.success:
            if self.result_type == "echarts":
                result["option"] = self.config
            else:
                result.update(self.config)
        else:
            result["error"] = self.error_message
        
        return result


@dataclass
class QueryContext:
    """查询上下文模型"""
    question: str
    language: str = "zh"
    user_preferences: Dict[str, Any] = field(default_factory=dict)
    additional_context: str = ""
    
    def __post_init__(self):
        """验证查询上下文"""
        if not self.question.strip():
            raise ValueError("问题不能为空")
        
        if self.language not in ["zh", "en"]:
            raise ValueError("语言必须是 'zh' 或 'en'")


@dataclass
class EngineConfig:
    """引擎配置模型"""
    llm_client: Any
    max_retries: int = 3
    timeout: int = 30
    enable_cache: bool = True
    cache_ttl: int = 3600
    safe_execution: bool = True
    max_data_rows: int = 10000
    debug_mode: bool = False
    
    def __post_init__(self):
        """验证引擎配置"""
        if self.max_retries < 0:
            raise ValueError("最大重试次数不能为负数")
        
        if self.timeout <= 0:
            raise ValueError("超时时间必须大于0")
        
        if self.cache_ttl <= 0:
            raise ValueError("缓存TTL必须大于0")
        
        if self.max_data_rows <= 0:
            raise ValueError("最大数据行数必须大于0")


@dataclass
class ErrorInfo:
    """错误信息模型"""
    error_type: str
    error_message: str
    error_code: str = ""
    traceback: str = ""
    timestamp: datetime = field(default_factory=datetime.now)
    context: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "error_type": self.error_type,
            "error_message": self.error_message,
            "error_code": self.error_code,
            "timestamp": self.timestamp.isoformat(),
            "context": self.context
        }


@dataclass
class VisualizationRequest:
    """可视化请求模型"""
    question: str
    data: Union[pd.DataFrame, Dict[str, pd.DataFrame]]
    chart_type: Optional[str] = None
    session_id: Optional[str] = None
    language: str = "zh"
    user_preferences: Dict[str, Any] = field(default_factory=dict)
    
    def __post_init__(self):
        """验证请求"""
        if not self.question.strip():
            raise ValueError("问题不能为空")
        
        if isinstance(self.data, pd.DataFrame) and self.data.empty:
            raise ValueError("数据不能为空")


@dataclass
class VisualizationResponse:
    """可视化响应模型"""
    success: bool
    chart_config: Optional[Dict[str, Any]]
    data: List[Dict[str, Any]]
    chart_type: Optional[str]
    execution_time: float
    session_id: str
    message: str = ""
    error: Optional[str] = None
    analysis_result: Optional[Dict[str, Any]] = None
    
    def __post_init__(self):
        """验证响应"""
        if not self.success and not self.error:
            raise ValueError("失败时必须提供错误信息")
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        result = {
            "success": self.success,
            "chart_type": self.chart_type,
            "execution_time": self.execution_time,
            "session_id": self.session_id,
            "message": self.message,
            "data": self.data
        }
        
        if self.success and self.chart_config:
            result["chart_config"] = self.chart_config
        
        if self.error:
            result["error"] = self.error
        
        if self.analysis_result:
            result["analysis_result"] = self.analysis_result
        
        return result